#coding=utf8

from numpy import *
import os

# classify using kNN
def kNNClassify(newInput, dataSet, labels, k):
    numSamples = dataSet.shape[0]  # shape[0] stands for the num of row

    ## step 1: calculate Euclidean distance
    # tile(A, reps): Construct an array by repeating A reps times
    # the following copy numSamples rows for dataSet
    diff = tile(newInput, (numSamples, 1)) - dataSet  # Subtract element-wise
    squaredDiff = diff ** 2  # squared for the subtract
    squaredDist = sum(squaredDiff, axis=1)  # sum is performed by row
    distance = squaredDist ** 0.5

    # argsort() returns the indices that would sort an array in a ascending order
    sortedDistIndices = argsort(distance)

    classCount = {}
    for i in xrange(k):
        voteLabel = labels[sortedDistIndices[i]]
        classCount[voteLabel] = classCount.get(voteLabel, 0) + 1

    maxCount = 0
    for key, value in classCount.items():
        if value > maxCount:
            maxCount = value
            maxIndex = key

    return maxIndex


# convert image to vector
def img2vector(filename):
    rows = 32
    cols = 32
    imgVector = zeros((1, rows * cols))
    fileIn = open(filename, 'r')
    for row in xrange(rows):
        lineStr = fileIn.readline()
        for col in xrange(cols):
            imgVector[0, row * 32 + col] = int(lineStr[col])

    return imgVector


# load dataSet
def loadDataSet():
    print "---Getting training set..."
    trainingFileList = os.listdir('digits/trainingDigits')  # load the training set
    numSamples = len(trainingFileList)

    trainX = zeros((numSamples, 1024))
    trainY = []
    for i in xrange(numSamples):
        filename = trainingFileList[i]
        trainX[i, :] = img2vector('digits/trainingDigits/' + filename)

        # get label from file name such as "1_18.txt"
        label = int(filename.split('_')[0])  # return 1
        trainY.append(label)

    print "---Getting testing set..."
    testingFileList = os.listdir('digits/testDigits')
    numSamples = len(testingFileList)
    testX = zeros((numSamples, 1024))
    testY = []
    for i in xrange(numSamples):
        filename = testingFileList[i]
        testX[i, :] = img2vector('digits/testDigits/%s' % filename)

        # get label from file name such as "1_18.txt"
        label = int(filename.split('_')[0])  # return 1
        testY.append(label)

    return trainX, trainY, testX, testY


# test hand writing class
def testHandWritingClass():
    print "step 1: load data..."
    trainDataset, trainLabel, testDataset, testLabel = loadDataSet()

    print "step 2: testing..."
    numTestSamples = testDataset.shape[0]
    matchCount = 0
    for i in xrange(numTestSamples):
        predict = kNNClassify(testDataset[i], trainDataset, trainLabel, 5)
        if predict == testLabel[i]:
            matchCount += 1
    accuracy = float(matchCount) / numTestSamples

    print "step 3: show the result..."
    print 'The classify accuracy is: %.2f%%' % (accuracy * 100)


testHandWritingClass()